package org.zjt.spark.book

import java.io.File

import org.apache.spark.mllib.classification.{NaiveBayes, NaiveBayesModel}
import org.apache.spark.mllib.linalg.SparseVector
import org.apache.spark.mllib.regression.LabeledPoint
import org.apache.spark.{SparkConf, SparkContext}

/**
  * 朴素贝叶斯 ： 分类算法
  *
  * 测试模型的准确性:1
  */
object MyNaiveBayes {

  val modelFile = "/Users/zhangjuntao/IdeaProjects/myproject/hw-bigdata/scala-demo/target/tmp/myNaiveBayesModel"
  val sourceFile = "/Users/zhangjuntao/IdeaProjects/myproject/hw-bigdata/scala-demo/src/main/resource/sample_libsvm_data.txt"
  val modelType = Array("multinomial", "bernoulli")

  def main(args: Array[String]): Unit = {
    val conf = new SparkConf().setAppName("NaiveBayesExample").setMaster("local[2]")
    val sc = new SparkContext(conf)
    val source = sc.textFile(sourceFile).coalesce(2).cache()
    val features = source.flatMap(_.split(" ").drop(1).map(_.split(":")(0))).distinct().cache()
    println(features.collect().mkString(","))
    println(features.count())
    val featureIndex = features.zipWithIndex().collectAsMap();
    //得到向量的索引
    val data = source.map {
      line => {
        val data = line.split(" ")
        val label = data.head.toInt
        val fields = data.tail

        //得到特征集合中的索引
        var feature2 = fields.map {
          field => {
            val feature = field.split(":")(0)
            val index = featureIndex.get(feature) match {
              case Some(a) => a.toInt
              case None => 0
            }
            val value = field.split(":")(1).toDouble
            (index, value)
          }
        }
        feature2 = feature2.sortBy(_._1)
        //println(feature2.mkString(","))
        val vector = new SparseVector(featureIndex.size, feature2.map(a => a._1), feature2.map(a => a._2))
        new LabeledPoint(label, vector)
      }
    }


    // Split data into training (60%) and test (40%).
    val Array(training, test) = data.randomSplit(Array(0.6, 0.4))
    println("training=%s\ttest=%s".format(training.count(), test.count()))
    val model = NaiveBayes.train(training, lambda = 1.0, modelType = modelType(0))
    val predictionAndLabel = test.map(p => (model.predict(p.features), p.label))


    //测试模型的准确性   测试模型的准确性:1.0
    val accuracy = 1.0 * predictionAndLabel.filter(x => x._1 == x._2).count() / test.count()
    println("测试模型的准确性:" + accuracy)

    val file = new File("/Users/zhangjuntao/IdeaProjects/myproject/hw-bigdata/scala-demo/target/tmp")
    if (file.exists())
      file.delete()

    // 保存模型
    model.save(sc, modelFile)


    //加载模型
    val sameModel: NaiveBayesModel = NaiveBayesModel.load(sc, modelFile)

    sc.stop()
  }
}
